# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import getpass
import logging
import os
import tempfile

import cloudpickle

from coin.util.slurm.base import (write_file, pickle_load_from_file)


class SlurmTask(object):
  def __init__(self,
               func,
               args=None,
               kwargs=None,
               compress=True,
               compresslevel=5,
               task_file_prefix=None,
               user=None,
               grid=None,
               shared_tempdir=None):
    assert callable(func)
    if args is None:
      args = []
    if kwargs is None:
      kwargs = {}
    self._pickled = cloudpickle.dumps((func, args, kwargs))

    self._compress = compress
    self._compresslevel = compresslevel

    self._task_file_prefix = 'slurm-tmp-'
    if task_file_prefix is not None:
      self._task_file_prefix = 'slurm-%s-' % str(task_file_prefix)

    user = user or getpass.getuser()
    grid = grid or 'iosg'
    shared_tempdir = shared_tempdir or '/remote/{grid}/home/{user}/scratch/tmp'
    self._shared_tempdir = shared_tempdir.format(grid=grid, user=user)

    self._input_file = None
    self._output_filepath = None

  def __enter__(self):
    if not os.path.isdir(self._shared_tempdir):
      os.makedirs(self._shared_tempdir, exist_ok=True)

    input_file_suffix = '-in' if not self._compress else '-in.gz'
    self._input_file = tempfile.NamedTemporaryFile(mode='w+b',
                                                   suffix=input_file_suffix,
                                                   prefix=self._task_file_prefix,
                                                   dir=self._shared_tempdir,
                                                   delete=True)

    output_file_suffix = '-out' if not self._compress else '-out.gz'
    self._output_filepath = '%s%s' % (self._input_file.name[:-len(input_file_suffix)],
                                      output_file_suffix)

    write_file(self._pickled,
               self._input_file,
               flush=True,
               compress=self._compress,
               compresslevel=self._compresslevel)

    input_file_size = os.stat(self.input_filepath).st_size
    if input_file_size > 2**20:  # 1MBytes
      logging.warning('Task input file size: %.2f MB', input_file_size / 2.**20)

    return self

  def __exit__(self, exc_type, exc_value, exc_traceback):
    self._input_file.close()
    if os.path.exists(self.output_filepath):
      os.remove(self.output_filepath)

  def read_output(self):
    result, exc_info = pickle_load_from_file(self.output_filepath)
    return result, exc_info

  def read_output_to_future(self, out_future):
    result, exc_info = self.read_output()
    if exc_info is None:
      out_future.set_result(result)
    else:
      try:
        out_future.set_exception_info(exc_info)
      except AttributeError:
        out_future.set_exception(exc_info[1])
    return result, exc_info

  @property
  def input_filepath(self):
    return self._input_file.name

  @property
  def output_filepath(self):
    return self._output_filepath

  def get_command_line(self, venv=None):
    args = ['coin/util/slurm/run_task.sh', self.input_filepath, self.output_filepath]
    if venv:
      args += [venv]
    return args
